package xyz.scootaloo.oj.offer

import com.sun.source.tree.Tree
import org.junit.Test
import xyz.scootaloo.console.app.support.Tester
import xyz.scootaloo.oj.ds.TreeNode

/**
 * @author flutterdash@qq.com
 * @since 2021/4/12 15:24
 */
class P54 {

    private fun kthLargest(root: TreeNode?, k: Int): Int {
        if (root == null) return 0
        val res = IntArray(1)
        dfs(root, k, 0, res)
        return res[0]
    }

    private fun dfs(root: TreeNode, k: Int, count: Int, res: IntArray): Int {
        var sum = count
        if (root.right != null) {
            sum = dfs(root.right, k, sum, res)
        }
        sum += 1
        if (sum == k) {
            res[0] = root.`val`
        }

        if (root.left != null) {
            sum = dfs(root.left, k, sum, res)
        }
        return sum
    }

    @Test
    fun test0() {
        val node1 = TreeNode(1)
        val node2 = TreeNode(2)
        val node3 = TreeNode(3)
        val node4 = TreeNode(4)

        node3.right = node4
        node3.left = node1
        node1.right = node2

        Tester.createTest(this::testHelper)
            .addCase(Wrapper(node3, 1), 4)
            .testAndShow()
    }

    @Test
    fun test1() {
        val node1 = TreeNode(1)
        val node2 = TreeNode(2)
        val node3 = TreeNode(3)
        val node4 = TreeNode(4)
        val node5 = TreeNode(5)
        val node6 = TreeNode(6)

        node5.right = node6
        node5.left = node3

        node3.left = node2
        node3.right = node4

        node2.left = node1

        Tester.createTest(this::testHelper)
            .addCase(Wrapper(node5, 3), 4)
            .addCase(Wrapper(node5, 4), 3)
            .testAndShow()
    }

    private fun testHelper(wrapper: Wrapper): Int {
        return kthLargest(wrapper.treeNode, wrapper.k)
    }

    inner class Wrapper(val treeNode: TreeNode, val k: Int)

}
